#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2011
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Author: Jesus Carrero <j.o.carrero@gmail.com>
#
# Look at this page for using diag sparse matrices in scipy:
# http://pod.tamu.edu/~lmkli/reference/python/scipy_html/generated/scipy.sparse.spdiags.html
#
# On Cubic Splines and the implementation here:
# Numerical Algorithms in C.
# [Gisela Engeln-Mullges and Frank Uhligi.] (1996) Chapter 10

from numpy import zeros, ones, exp, linspace
from numpy import array, r_, c_, roll
from scipy.sparse.construct import spdiags
from scipy.sparse.linalg import spsolve as spsolve


class Spline:

    """
        This is an implementation of natural cubic splines
        Sp(x) = a + b * (x-xj) + c *(x-xj)^2 + d *(x-xj)^3
    """

    __slots__ = ['m_xi', 'm_yi', 'm_spc']

    def __init__(self, _xi=None, _yi=None):
        self.m_xi = _xi
        self.m_yi = _yi
        self.m_spc = None
        if not (_xi is None and _yi is None):
            self.spline(_xi, _yi)

    def _compSplineCoeff(self):
        (x, y) = (self.m_xi, self.m_yi)

        one = ones((1, 1))
        zer = zeros((1, 1))
        dx = roll(x, -1) - x
        dx = dx[:-1]  # INTERVAL LENGTH
        d = 2 * (dx + roll(dx, -1))
        d = r_[r_[one, d[1:]], one]
        u = r_[zer, dx.copy()]
        u[1] = 0
        l = r_[dx.copy(), zer]
        l[-2] = 0

        b = (roll(y, -1) - y)[:-1] / dx
        b = 3. * (roll(b, -1) - b)
        b = r_[r_[zer, b[:-1]], zer]

        data = array([list(d.flatten()), list(l.flatten()),
                     list(u.flatten())])

        diags = array([0, -1, 1])
        M = spdiags(data, diags, d.size, d.size).tocsc()
        c = spsolve(M, b, use_umfpack=False)
        c = c.reshape((c.size, 1))

        # SPLINE COEFFICIENTS

        spc = zeros((x.size - 1, 4), 'float64')
        spc[:, 0] = y[:-1].flatten()  # a
        spc[:, 1] = ((roll(y, -1) - y)[:-1] / dx - dx * (roll(c, -1)
                     + 2 * c)[:-1] / 3.).flatten()  # b
        spc[:, 2] = c[:-1].flatten()  # c
        spc[:, 3] = 1. / 3. * ((roll(c, -1) - c)[:-1] / dx).flatten()  # d
        self.m_spc = spc

    def get_grid(self):
        return self.m_xi

    def getFuncValue(self):
        return self.m_yi

    def spline(self, _xi, _yi):
        self.m_xi = _xi.reshape((_xi.size, 1))
        self.m_yi = _yi.reshape((_yi.size, 1))
        self._compSplineCoeff()

    def value(self, x, i):
        a = self.m_spc[i, :]
        dx = x - self.m_xi[i]
        return a[0] + dx * (a[1] + dx * (a[2] + dx * a[3]))

    def derivative(self, x, i):
        a = self.m_spc[i, :]
        dx = x - self.m_xi[i]
        return a[1] + dx * (2. * a[2] + dx * 3. * a[3])

    def secondDerivative(self, x, i):
        a = self.m_spc[i, :]
        dx = x - self.m_xi[i]
        return 2. * a[2] + 6. * a[3] * dx


def f(x):
    return 10 + 20 * x - 30 * x ** 2
    return x ** 3
    return ones(x.shape)


if __name__ == '__main__':
    nSeg = 100
    x = linspace(0, nSeg, nSeg, False)
    y = f(x)
    mySpline = Spline(x, y)

    # chk Results

    shift = 0.5
    yt = f(x + shift * ones(x.shape))
    spVals = array([mySpline.value([x[s] + shift], s) for s in
                   range(x.size - 1)]).flatten()
    print ' x, error '
    print c_[(x + shift)[:-1], yt[:-1] - spVals]

